[6403893][ONNX][AutoCast] Fix false-success TRT parse for large ModelProto#1928
Conversation
|
/claude review |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughAdds oversized-model detection for in-memory ONNX models, routes applicable models through temporary external-data files for TensorRT parsing, and tests parity across file-path, in-memory, and forced file-backed modes. ChangesFile-backed ONNX parsing support
Estimated code review effort: 2 (Simple) | ~15 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant get_custom_layers
participant _requires_file_backed_parse
participant TempFileSave
participant TrtParser
Caller->>get_custom_layers: onnx_path (str or ModelProto)
alt onnx_path is string
get_custom_layers->>TrtParser: parse_from_file(onnx_path)
else onnx_path is ModelProto
get_custom_layers->>_requires_file_backed_parse: check ByteSize()
alt requires file-backed parse
_requires_file_backed_parse-->>get_custom_layers: True
get_custom_layers->>TempFileSave: deepcopy + save external data
get_custom_layers->>TrtParser: parse_from_file(temp_path)
else
_requires_file_backed_parse-->>get_custom_layers: False
get_custom_layers->>TrtParser: parse(SerializeToString())
end
end
TrtParser-->>get_custom_layers: parsed layers or errors
Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| with tempfile.TemporaryDirectory(prefix="modelopt_trt_") as tmpdir: | ||
| model_path = os.path.join(tmpdir, "model.onnx") | ||
| # Deep-copy so externalization doesn't strip weights from the caller's model. | ||
| model_copy = copy.deepcopy(onnx_path) |
There was a problem hiding this comment.
[SUGGESTION] copy.deepcopy(onnx_path) transiently doubles peak RAM for exactly the models this path targets — an 8GB ModelProto needs ~16GB before onnx.save streams weights to disk. The deep-copy is necessary for correctness (onnx.save(save_as_external_data=True) mutates the proto in place, stripping raw_data from the caller's model), so this isn't a bug — but on the 8GB VLA trunk case cited in the PR description it may be tight on constrained hosts.
If it ever bites, an alternative that avoids the copy is to externalize in place and restore afterward (e.g. onnx.save(onnx_path, ...) then onnx.load_external_data_for_model(onnx_path, tmpdir) to re-hydrate the caller's proto), or wrap in try/finally to guarantee re-hydration on parse failure. Non-blocking; the current approach is the safest correctness-wise.
There was a problem hiding this comment.
The peak memory usage is at ReferenceRunner where we define an auxiliary graph with every intermediate tensor being a graph output. The 2x memory usage here is not a bottleneck.
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Scope: Reviewed the 2 files in this PR's authoritative changed-file list (modelopt/onnx/trt_utils.py, tests/gpu/onnx/quantization/test_plugin.py). The two-dot git diff origin/main HEAD surfaced many unrelated files due to the shallow-checkout divergence between the local origin/main tip and the true merge base — those are not part of this PR and were excluded.
Findings: CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1
Assessment — low risk. Focused, well-targeted bugfix. get_custom_layers previously silently returned an empty network (0 layers/0 tensors) for in-memory ModelProtos at/above the protobuf 2 GiB limit because TensorRT's parser.parse(bytes) fails quietly while SerializeToString() succeeds. The fix routes over-limit (or un-sizeable) in-memory models through a temporary external-data file + parse_from_file(), matching the existing string-path branch and ModelOpt's file-backed policy elsewhere in this module.
Verified:
onnx.checker.MAXIMUM_PROTOBUFis a valid constant, already used with the same>= MAXIMUM_PROTOBUFidiom at trt_utils.py:382.- The
ByteSize()-overflowexceptbranch conservatively routes to the file path, consistent withload_onnx_model's handling. copy.deepcopycorrectly shields the caller's model from in-place externalization byonnx.save; both in-memory callers pass a model the caller still owns.- The regression test forces the file-backed path via monkeypatch and cross-checks results against both the in-memory fast path and the string-path baseline. Backward compatible — the small-model fast path is unchanged.
The single SUGGESTION concerns transient peak-RAM doubling from the deep-copy on multi-GB models — non-blocking, and the current approach is the correctness-safe choice.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1928 +/- ##
==========================================
+ Coverage 76.50% 77.25% +0.74%
==========================================
Files 525 525
Lines 59257 59273 +16
==========================================
+ Hits 45337 45790 +453
+ Misses 13920 13483 -437
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
gcunhase
left a comment
gcunhase
left a comment
There was a problem hiding this comment.
Do you think the comment from Claude review should be a concern?
The single SUGGESTION concerns transient peak-RAM doubling from the deep-copy on multi-GB models.
Otherwise, LGTM!
@gcunhase Thanks for holding me accountable... Replied in the original thread. |
b52dc3f to
c8a5a30
Compare
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/gpu/onnx/quantization/test_plugin.py`:
- Line 163: Move the `modelopt.onnx.trt_utils as trt_utils` import out of the
test body and into the module-level imports in `test_plugin.py`. This import is
only used to support `monkeypatch.setattr`, so it should be declared alongside
the other top-level `modelopt.onnx.trt_utils` imports rather than inside the
test. Remove the in-function import and keep the `trt_utils` reference available
for the existing monkeypatch usage.
🪄 Autofix (Beta)
❌ Autofix failed (check again to retry)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e2d67f9b-19fd-4e4f-a16a-79d54c790fd8
📒 Files selected for processing (2)
modelopt/onnx/trt_utils.pytests/gpu/onnx/quantization/test_plugin.py
|
Note Autofix is a beta feature. Expect some limitations and changes as we gather feedback and continue to improve it. An unexpected error occurred while generating fixes: Not Found - https://docs.github.com/rest/git/refs#get-a-reference |
d5228a7 to
55b0296
Compare
55b0296 to
f764597
Compare
…Proto get_custom_layers silently returned 0 layers/tensors for in-memory ModelProtos at/above the protobuf 2GiB limit. Route such models through a temporary external-data file and parse_from_file(), matching the string-path behavior. Verified on the 8.12GB VLA trunk (28 layers/5468 tensors) and a synthetic 2.2GB model. Adds a regression test forcing the file-backed path. Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com>
Signed-off-by: Gal Hubara-Agam <96368689+galagam@users.noreply.github.com> Signed-off-by: Gal Hubara Agam <96368689+galagam@users.noreply.github.com>
f764597 to
38f9c7d
Compare
2 participants
What does this PR do?
Type of change: ? Bugfix
get_custom_layers silently returned 0 layers/tensors for in-memory ModelProtos at/above the protobuf 2GiB limit. Route such models through a temporary external-data file and parse_from_file(), matching the string-path behavior. Verified on the 8.12GB VLA trunk (28 layers/5468 tensors) and a synthetic 2.2GB model. Adds a regression test forcing the file-backed path.
Usage
# Add a code snippet demonstrating how to use thisTesting
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: N/AAdditional Information
Summary by CodeRabbit
Summary
Bug Fixes
Tests
get_custom_layersreturns matching custom-layer lists and tensor metadata for in-memory versus file-path parsing, including a forced file-backed scenario.